package middleware

import (
	"fmt"
	"ginstudydemo.com/main/service"
	"github.com/gin-gonic/gin"
	"github.com/redis/go-redis/v9"
	"time"
)

type ViewRateLimiter struct {
	rdb *redis.Client
}

func NewViewRateLimiter(rdb *redis.Client) *ViewRateLimiter {
	return &ViewRateLimiter{rdb: rdb}
}

// RateLimit 增强版限流中间件
func (v *ViewRateLimiter) RateLimit(limit int, window time.Duration) gin.HandlerFunc {
	return func(c *gin.Context) {
		resourceID := c.Param("resourceId")
		ip := c.ClientIP()

		// 获取用户ID（根据实际系统调整）
		userID, _ := c.Get("userID")
		identifier := fmt.Sprintf("%s|%v|%s", ip, userID, resourceID)

		key := "view_rate:" + identifier

		// 使用Lua脚本保证原子操作
		luaScript := `
local current = redis.call('INCR', KEYS[1])
if current == 1 then
    redis.call('EXPIRE', KEYS[1], ARGV[1])
end
return current`

		result, err := v.rdb.Eval(c, luaScript, []string{key}, window.Seconds()).Int()
		if err != nil {
			service.ReturnErr(c, "9999", "系统繁忙")
			c.Abort()
			return
		}

		if result > limit {
			service.ReturnErr(c, "9999", "操作过于频繁")
			c.Abort()
			return
		}

		c.Next()
	}
}
